[None][feat] MTP one-model advanced_sampling_mode: skip redundant top-k / top-p filter kernels with additional config enum#16561
Conversation
0e79ccf to
4ef515a
Compare
jhaotingc
changed the title
advanced_sampling_mode: skip redundant top-k / top-p filter kernels with additional config enum
b96248c to
68a00f4
Compare
jhaotingc
added
the
api-compatible
68a00f4 to
6c7f810
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds deploy-time advanced sampling modes for MTP one-model speculative decoding, propagates the selected mode through metadata and sampling paths, conditionally skips top-k/top-p filtering, and adds validation tests and documentation. ChangesAdvanced sampling modes
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant MTPDecodingConfig
participant SpecMetadata
participant SpeculativeInterface
participant sampling_utils
participant FlashInfer
MTPDecodingConfig->>SpecMetadata: configure advanced_sampling_mode
SpecMetadata->>SpeculativeInterface: provide mode for speculative sampling
SpeculativeInterface->>sampling_utils: resolve effective top_k and top_p
sampling_utils-->>SpeculativeInterface: return filters or None
SpeculativeInterface->>sampling_utils: sample draft tokens or compute rejection probabilities
sampling_utils->>FlashInfer: apply enabled sampling kernels
FlashInfer-->>sampling_utils: return probabilities or token IDs
sampling_utils-->>SpeculativeInterface: return sampling results
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py (2)
39-53: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise rejection-disabled behavior explicitly.
Lines 41 and 49 rely on the default for
use_rejection_sampling; set it toFalseso the tests directly prove both sides of the contract.Proposed test adjustment
- cfg = MTPDecodingConfig(max_draft_len=1, advanced_sampling_mode="no_topk") + cfg = MTPDecodingConfig( + max_draft_len=1, + advanced_sampling_mode="no_topk", + use_rejection_sampling=False, + ) ... - MTPDecodingConfig(max_draft_len=1, advanced_sampling_mode=mode) + MTPDecodingConfig( + max_draft_len=1, + advanced_sampling_mode=mode, + use_rejection_sampling=False, + )As per path instructions, coverage in
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.pyneeds direct validation of the configured behavior.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py` around lines 39 - 53, Update both MTPDecodingConfig constructions in test_no_topk_does_not_require_rejection and test_topp_disabling_modes_require_rejection to explicitly pass use_rejection_sampling=False for the rejection-disabled cases, while keeping the existing rejection-enabled configuration unchanged.Source: Path instructions
90-125: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover NO_TOPK rejection of greedy requests.
This only tests temperature
0.7. Add a case at or belowGREEDY_TEMPERATURE_THRESHOLDthat asserts the NO_TOPK request-validation path rejects it, plus a just-above-threshold success case.Based on PR objectives, NO_TOPK disallows greedy requests. As per path instructions, coverage in
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.pyis insufficient for that contract.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py` around lines 90 - 125, Extend test_no_topk_matches_full_bit_for_bit or add a focused test covering NO_TOPK temperature validation. Use temperatures at or below GREEDY_TEMPERATURE_THRESHOLD to assert the NO_TOPK request is rejected, and just above the threshold to assert it succeeds. Reuse sampling_batch_spec_dec_one_model and the existing test setup symbols, while preserving the current bit-for-bit comparison coverage.Source: Path instructions
tensorrt_llm/llmapi/llm_args.py (1)
2484-2504: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the validator return annotation.
_log_advanced_sampling_modeshould declare-> "MTPDecodingConfig".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/llmapi/llm_args.py` around lines 2484 - 2504, The _log_advanced_sampling_mode validator is missing its return annotation. Update that method signature to declare -> "MTPDecodingConfig", preserving its existing validation logic and return self behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py`:
- Around line 890-897: Update the advanced_sampling_mode branch in the direct
sampling path to use compute_probs_from_logits with the selected mode, then
sample via sampling_from_probs_op instead of always calling
top_p_sampling_from_probs_op. Ensure both NO_TOPP and NO_TOPK_NO_TOPP bypass
top_p filtering while preserving the existing temperature, seed, and offset
behavior.
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 830-846: Initialize and update self.has_greedy_requests while
scanning per-request sampling parameters, ensuring it reflects any greedy row in
mixed batches. Apply the existing mixed-batch rejection to every non-FULL
advanced sampling mode, including no_topk and no_topk_no_topp, while preserving
the all-greedy and warmup-dummy exemptions.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 2475-2482: Validate advanced_sampling_mode after the final
spec-decoding mode is determined, accepting non-FULL values only for the
MTP-Eagle one-model branch. Reject unsupported non-FULL modes for vanilla and
two-model MTP instead of silently ignoring them, while preserving FULL as the
default and existing metadata behavior.
In
`@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py`:
- Line 30: Add complete type annotations to every function in this test module,
including all test functions and the spy helper: annotate each parameter with
its precise type and use -> None for functions that do not return a value.
Update the functions represented by test_advanced_sampling_mode_enum_values and
the other referenced test/spy definitions without changing their behavior.
---
Nitpick comments:
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 2484-2504: The _log_advanced_sampling_mode validator is missing
its return annotation. Update that method signature to declare ->
"MTPDecodingConfig", preserving its existing validation logic and return self
behavior.
In
`@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py`:
- Around line 39-53: Update both MTPDecodingConfig constructions in
test_no_topk_does_not_require_rejection and
test_topp_disabling_modes_require_rejection to explicitly pass
use_rejection_sampling=False for the rejection-disabled cases, while keeping the
existing rejection-enabled configuration unchanged.
- Around line 90-125: Extend test_no_topk_matches_full_bit_for_bit or add a
focused test covering NO_TOPK temperature validation. Use temperatures at or
below GREEDY_TEMPERATURE_THRESHOLD to assert the NO_TOPK request is rejected,
and just above the threshold to assert it succeeds. Reuse
sampling_batch_spec_dec_one_model and the existing test setup symbols, while
preserving the current bit-for-bit comparison coverage.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 87e95728-f8e1-4129-8afa-df75e9e6ff80
📒 Files selected for processing (5)
tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py
| from tensorrt_llm.llmapi.llm_args import AdvancedSamplingMode, MTPDecodingConfig | ||
|
|
||
|
|
||
| def test_advanced_sampling_mode_enum_values(): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add required annotations to every function.
All test functions and spy lack parameter and return annotations. Add precise parameter types and -> None for tests.
As per coding guidelines, “Annotate every function.”
Also applies to: 39-39, 46-46, 65-67, 71-74, 90-90
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py`
at line 30, Add complete type annotations to every function in this test module,
including all test functions and the spy helper: annotate each parameter with
its precise type and use -> None for functions that do not return a value.
Update the functions represented by test_advanced_sampling_mode_enum_values and
the other referenced test/spy definitions without changing their behavior.
Source: Coding guidelines
6c7f810 to
c8eb118
Compare
| # no-op at k=vocab) so tokens match "full" bit-for-bit, avoiding any RNG discrepancy. | ||
| logits = vanilla.safely_apply_temperature_inplace(logits, temperatures) | ||
| probs = torch.softmax(logits, dim=-1) | ||
| return top_p_sampling_from_probs_op(probs, top_p, seed=seed, offset=offset) |
There was a problem hiding this comment.
Please following compute_probs_from_logits_op, add a sample_from_logits_op; resolve advanced_sampling_mode once on the speculative side and call it directly, removing the sampling_batch_spec_dec_one_model wrapper layer.
@torch.compile(options={"max-autotune": True})
def sample_from_logits_op(
logits: torch.Tensor,
temperatures: torch.Tensor,
top_k: Optional[torch.Tensor] = None,
top_p: Optional[torch.Tensor] = None,
seed: Optional[torch.Tensor] = None,
offset: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if top_k is not None:
logits = flashinfer.sampling.top_k_mask_logits(logits, top_k)
probs = flashinfer.sampling.softmax(
logits, temperatures, enable_pdl=get_env_enable_pdl())
if top_p is not None:
return flashinfer.sampling.top_p_sampling_from_probs(
probs, top_p, seed=seed, offset=offset)
return flashinfer.sampling.sampling_from_probs(probs, seed=seed, offset=offset)
There was a problem hiding this comment.
@zhaoyangwang-nvidia I was having concern the minute diff when user turn on no_topp for non-rejection path.
In your suggestion, if user decide to turn on no_topp and pass in top_p 1.0, the kernel will be SamplingFromProbKernel instead of TopPSamplingFromProbKernel. The two kernels handle sorting differently.
Before #15775 where it removes drafter sampling for default case, the diff is more significant. So I disabled no_topp use case when rejection sampling False to make it consistent. The no_rejection sampling path is already having TopP kernel, switching to sampling from prob does not help much because 1 kernel still needs to run. The rejection path already has TopK/TopP added by you so I reused the same logic for no_topk_no_topp.
If you're ok with the behavior changes for turning on no_topp in no_rejection path, now that draft step sampling is removed, I'm ok with it as well.
There was a problem hiding this comment.
Also, another reason of not able to discard this function is this line. This line makes ALL full path having one GREEDY result calculated and one sampled result calculated, just to make sure we take care of when temperatures <= vanilla.GREEDY_TEMPERATURE_THRESHOLD
Using this function, you suggested changes the behavior completely, and the same issue isn't considered in #15775
There was a problem hiding this comment.
#15542 seems to add this greedy & advanced together, is this intended? should this be kept?
There was a problem hiding this comment.
hmm ok didn't see the comments here. I'll change the behavior, thanks.
There was a problem hiding this comment.
Sorry for the late reply — timezone difference on my side.
Both behavior differences you raised are valid observations:
SamplingFromProbKernel vs TopPSamplingFromProbKernel producing different tokens for the same distribution when no_topp is enabled;
the greedy rows no longer going through an explicit argmax + torch.where selection, but instead being handled natively via the DISABLE_TEMP_VAL sentinel temperature.
That said, I'd argue both are correct behaviors that differ only by floating-point accumulation order / kernel implementation details, not accuracy bugs. In case 1, top_p=1.0 means no filtering, so any properly implemented sampling kernel is a valid realization of the same distribution. In case 2, softmax(logits / DISABLE_TEMP_VAL) collapses to a one-hot at the argmax, so sampling from it returns the argmax token — the greedy semantics are preserved.
Do you have any real cases where these differences caused an actual accuracy problem (e.g., a measurable eval-score regression, or a rejection-rate change)? If so, let's discuss further and I'm happy to reconsider. If not, I'd prefer to keep the code simple and unified — one mode-agnostic sampling op with the filters resolved once on the speculative side — rather than maintaining separate paths just for bit-for-bit token equivalence, which we don't guarantee across versions anyway.
There was a problem hiding this comment.
Thank you! Changed to what's suggested on comment, will click resolve before merge.
| # would be sampled instead of argmax'd. All-greedy batches use the greedy graph; warmup | ||
| # dummies are exempted above. | ||
| if (self.advanced_sampling_mode != "full" | ||
| and not self.is_all_greedy_sample and self.has_greedy_requests): |
There was a problem hiding this comment.
Please following how the rejection path handles greedy rows when computing draft/target probs (the sentinel-temperature approach — greedy rows carry DISABLE_TEMP_VAL, so softmax(logits / DISABLE_TEMP_VAL) collapses to a one-hot at the argmax and sampling from it returns the argmax token), the no_topk fast path could support greedy rows natively as well, which would eliminate this mixed-batch check and error.
There was a problem hiding this comment.
+1, no need to ValueError here
There was a problem hiding this comment.
remove all constraint, now the no_topk/no_topp will apply to no rejection sampling (the target sampling only) if pre-set.
There was a problem hiding this comment.
Launching sweeps to verify perf.
| "Token ID marking end of thinking phase. Strict acceptance resumes after this." | ||
| ) | ||
|
|
||
| advanced_sampling_mode: AdvancedSamplingMode = Field( |
There was a problem hiding this comment.
Please update document in docs/source/features/sampling.md
| # Whether to use rejection sampling for one-model speculative decoding. | ||
| use_rejection_sampling: bool = False | ||
| # Advanced-sampling specialization (deploy-time; from MTPDecodingConfig.advanced_sampling_mode). | ||
| advanced_sampling_mode: str = "full" |
There was a problem hiding this comment.
AdvancedSamplingMode is defined as an enum in llm_args.py, but everything downstream degrades to bare strings, and the "which modes skip top_k" predicate in ("no_topk", "no_topk_no_topp") is duplicated across sampling_utils.py (twice) and this file with nothing keeping the literals in sync as the enum evolves. Consider a single source of truth — e.g. properties on the enum (mode.skips_top_k / mode.skips_top_p) or derived boolean fields on SpecMetadata — so consumers read a boolean instead of each parsing the string. (non-blocking)
There was a problem hiding this comment.
Made all field to be enum instead of str, only pass in str in extra-llm-api-config.yml from user.
|
Hi @mikeiovine Please take a look at this PR — it adds advanced_sampling_mode to skip the top-k / top-p filter kernels for temperature-only deployments. Any comments on this implementation approach? |
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
I think the optimization makes sense but I'm a bit worried about the API. As you point out in the description, we could make it more usable by capturing more graphs, but that's going to make the warmup too long and is not sustainable.
At least the default is FULL. This should be considered an optimization for advanced use cases IMO
| "Token ID marking end of thinking phase. Strict acceptance resumes after this." | ||
| ) | ||
|
|
||
| advanced_sampling_mode: AdvancedSamplingMode = Field( |
There was a problem hiding this comment.
Should be part of the base config
There was a problem hiding this comment.
Moved to base config
| and not self.speculative_config.spec_dec_mode. | ||
| is_mtp_eagle_one_model()): |
There was a problem hiding this comment.
You can delete this. I am removing MTP one model mode
| # would be sampled instead of argmax'd. All-greedy batches use the greedy graph; warmup | ||
| # dummies are exempted above. | ||
| if (self.advanced_sampling_mode != "full" | ||
| and not self.is_all_greedy_sample and self.has_greedy_requests): |
There was a problem hiding this comment.
+1, no need to ValueError here
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
f2479d0 to
729831c
Compare
|
Since this adds a new field to DecodingBaseConfig, per the repo policy for LLM args / nested-config changes, please run python3 scripts/generate_llm_args_golden_manifest.py and commit the updated tensorrt_llm/usage/llm_args_golden_manifest.json — the manifest currently doesn't include advanced_sampling_mode, which will fail CI. Also note that new config fields require approval from the telemetry/privacy CODEOWNERs. |
|
@NVIDIA/trt-llm-doc-owners Please help to review this PR, thanks~ |
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
729831c to
8265a20
Compare
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
8265a20 to
7c56183
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py (2)
31-31: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse a snake_case local variable.
Mis a local variable; rename it to a descriptive snake_case name.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py` at line 31, Rename the local variable M assigned to AdvancedSamplingMode to a descriptive snake_case name, and update all references to it within the relevant test scope.Source: Coding guidelines
143-155: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftExercise the speculative paths claimed by this test.
This only constructs one
MTPDecodingConfig; it does not invoke draft, rejection, block, batch, or dynamic-tree sampling. Add path-level coverage, or narrow the test name/docstring to config construction.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py` around lines 143 - 155, Expand test_advanced_mode_accepted_on_all_spec_paths to exercise each claimed speculative sampling path, including draft, rejection, block, batch, and dynamic-tree handling, using the configured advanced_sampling_mode. If those paths cannot be covered here, rename the test and revise its docstring to describe only MTPDecodingConfig construction rather than all speculative paths.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In
`@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py`:
- Line 31: Rename the local variable M assigned to AdvancedSamplingMode to a
descriptive snake_case name, and update all references to it within the relevant
test scope.
- Around line 143-155: Expand test_advanced_mode_accepted_on_all_spec_paths to
exercise each claimed speculative sampling path, including draft, rejection,
block, batch, and dynamic-tree handling, using the configured
advanced_sampling_mode. If those paths cannot be covered here, rename the test
and revise its docstring to describe only MTPDecodingConfig construction rather
than all speculative paths.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cb0870b5-8034-46fd-bb19-e2c4118be7c3
📒 Files selected for processing (8)
docs/source/features/sampling.mdtensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.pytensorrt_llm/_torch/speculative/eagle3_dynamic_tree.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py
🚧 Files skipped from review as they are similar to previous changes (7)
- docs/source/features/sampling.md
- tensorrt_llm/_torch/speculative/utils.py
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
- tensorrt_llm/_torch/speculative/interface.py
|
PR_Github #61082 [ run ] triggered by Bot. Commit: |
|
PR_Github #61082 [ run ] completed with state
|
3d53d8a to
4e1349c
Compare
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
|
/bot run |
|
PR_Github #61400 [ run ] triggered by Bot. Commit: |
|
PR_Github #61400 [ run ] completed with state
|
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
4e1349c to
a49a3cb
Compare
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
|
|
/bot run --disable-fail-fast |
|
PR_Github #61443 [ run ] triggered by Bot. Commit: |
|
PR_Github #61443 [ run ] completed with state
|
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
d831d6b to
00dd072
Compare
|
Hi @Tabrizian @zhaoyangwang-nvidia would you mind review the sampling_util one model sampling there again? due to #16590 Thank you! |
00dd072 to
57c90ff
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61626 [ run ] triggered by Bot. Commit: |
|
PR_Github #61626 [ run ] completed with state
|
…els in the one-model speculative sampler Add a deploy-time `advanced_sampling_mode` enum on `DecodingBaseConfig` (default FULL) that lets the one-model speculative sampler skip flashinfer's top-k mask and top-p renorm kernels when the deployment disables those filters. When a skipped filter is already disabled the output matches FULL bit-for-bit, so this is a lossless throughput optimization for fixed-config / temperature-only deployments (measured ~8-28% OTPS/gpu gain on Qwen3.6-35B-A3B-NVFP4, B200, MTP draft_len=3). - Enum end-to-end: `AdvancedSamplingMode.skips_top_k` / `.skips_top_p` are the single source of truth; pydantic parses the config string into the enum, `SpecMetadata` carries it, and `resolve_advanced_sampling_filters` resolves it once on the speculative side so the ops stay mode-agnostic. - Unified sampler: `sample_from_logits_op` (top_k mask -> softmax -> top_p sample / plain sample) replaces the `sampling_batch_spec_dec_one_model` wrapper; a `None` top_k / top_p skips that filter's kernel. - Greedy rows are handled natively via the `DISABLE_TEMP_VAL` sentinel temperature (softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches with no special-case guard. - All modes work with `use_rejection_sampling` on or off; the two are independent. - Regenerated `tensorrt_llm/usage/llm_args_golden_manifest.json` (adds `speculative_config.advanced_sampling_mode`). - Unit tests in `tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py` and docs in `docs/source/features/sampling.md`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
57c90ff to
64cee60
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61672 [ run ] triggered by Bot. Commit: |
|
PR_Github #61672 [ run ] completed with state
|
4 participants
Description
For one-model speculative decoding (MTP-Eagle one-model), the advanced sampler runs
flashinfer's top-k mask (
RadixTopKMaskLogits) and top-p renorm (AirTopPRenorm*)kernels on every draft layer and the target verification — even when the
deployment's requests disable those filters (
top_k=0,top_p=1). In that case thefilters are mathematical no-ops (top-k mask at
k=vocab, top-p renorm atp=1), butthe kernels still launch: measured at ~290 µs (top-k mask) + ~1.34 ms (top-p
renorm) per decode step at
draft_len=3(nsys, Qwen3.6-35B-A3B-NVFP4, B200). Forfixed-config / temperature-only deployments this is pure overhead.
Add a deploy-time
advanced_sampling_modeenum onDecodingBaseConfig(availableto any speculative config; default
FULL) that tells the one-model sampler whichfilters the deployment disables, so their flashinfer kernels are skipped:
FULL(default)NO_TOPKNO_TOPPNO_TOPK_NO_TOPPDesign:
AdvancedSamplingMode.skips_top_k/.skips_top_pproperties are the only place the mode→filter mapping lives. The enum is carried
end-to-end: pydantic parses the YAML/JSON string into the enum on the config,
SpecMetadatastores the enum, and there are no bare-string predicates downstream.resolve_advanced_sampling_filters(mode, top_k, top_p)returns effective tensors with a disabled filter set toNone, so thesampler/prob ops stay mode-agnostic.
sample_from_logits_op(logits, temperatures, top_k, top_p, seed, offset)(mirroringcompute_probs_from_logits_op) is the single one-model sampler:a
Nonetop_kskips the mask kernel; aNonetop_psamples viasampling_from_probsinstead oftop_p_sampling_from_probs. The rejection path'scompute_probs_from_logitsskips the top-k mask / top-p renorm identically.(
DISABLE_TEMP_VAL) set in_scan_one_model_sampling, sosoftmax(logits / DISABLE_TEMP_VAL)collapses to a one-hot at the argmax and sampling returns theargmax token. Every mode therefore supports mixed greedy + sampling batches with no
special-case error.
advanced_sampling_modeis a deploy-time setting and is notpart of the CUDA-graph key.
NO_TOPKwithtop_k=0), the probs — and therefore the sampled tokens and rejection acceptance —are bit-for-bit identical to
FULL; only the kernel launch is removed.use_rejection_samplingon or off; the two are independent.Benchmark
Qwen3.6-35B-A3B-NVFP4, TRT-LLM (1×TP1, MTP
draft_len=3), B200, rejection sampling ON,SPEED-Bench low-entropy (ISL≈2197, OSL=1024, temp=0.7), prefix-cache OFF. Values are
OTPS/gpu (output tok/s/gpu). Plot:
plots/pareto_rejection_conc.png.* Speedup =
NO_TOPK_NO_TOPP/FULL (temp only)— same sampling config(
top_k=0, top_p=1), isolating the effect of skipping the two filter kernels.Acceptance length is unchanged across all rows, confirming identical sampling behavior.
FULLwith realtop_k=20 / top_p=0.95is shown as a reference for the actualfiltered-sampling throughput.
Result: for a temperature-only deployment,
NO_TOPK_NO_TOPPgives a consistent~8–28% throughput gain (≈+12% at high concurrency) at zero accuracy cost, by
removing the two redundant per-draft-layer filter kernels.
Test coverage
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py(12 cases, all passing on B200):
{full, no_topk, no_topp, no_topk_no_topp}and theskips_top_k/skips_top_pproperties; the field lives onDecodingBaseConfig(default
FULL); all four modes construct with or withoutuse_rejection_sampling(no gating).
resolve_advanced_sampling_filterssets the disabled filter toNoneper mode (full→neither,no_topk→top_k,no_topp→top_p,no_topk_no_topp→both) and passes kept filters through unchanged.NO_TOPKviasample_from_logits_opyields theexact same tokens as
FULLwhen top_k is disabled, for the same seed/offset, acrosstop_p ∈ {1.0, 0.9}(the top_k mask is a no-op atk=vocab).returns its argmax token under
NO_TOPKandNO_TOPK_NO_TOPP, validating theguard-free design.
FULLmodes construct on any spec-decoding path(config-time acceptance; unsupported paths fall back to
FULLbehavior at runtime).Telemetry manifest / CODEOWNERS
Because this adds a new
DecodingBaseConfigfield,tensorrt_llm/usage/llm_args_golden_manifest.jsonis regenerated via
scripts/generate_llm_args_golden_manifest.py(addsspeculative_config.advanced_sampling_mode,kind: categorical). The manifest--checkCI gate passes. Touching the usage-telemetry path requires approval from@NVIDIA/trt-llm-usage-telemetry-devs,@NVIDIA/trt-llm-oss-compliance, and@NVIDIA/trt-llm-noncommitted-api-review-committee. The new field is additive withdefault
FULL(API-compatible).Note — alternative design
Instead of a deploy-time setting that changes which sampler the single advanced CUDA
graph captures, each mode could be a separately captured CUDA graph selected at
replay by a graph key — the same mechanism
is_all_greedy_samplealready uses to switchbetween the greedy (argmax) and advanced-sampling graphs. That would allow per-batch
mode switching, but it multiplies the number of captured decode graphs (one set per
mode) and lengthens warmup. We chose the deploy-time enum (0 extra graphs) because the
sampling config is fixed per deployment, which is the target use case.
Test Coverage
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py(11 cases, all passing on B200):
{full, no_topk, no_topp, no_topk_no_topp};NO_TOPKis accepted without rejection sampling;NO_TOPPandNO_TOPK_NO_TOPPraise aValidationErrorunlessuse_rejection_sampling=True(and are accepted with it).
compute_probs_from_logitspassestop_k=Noneand/ortop_p=Noneinto the flashinfer op (checked with a spy),so the right kernel is skipped:
full→neither,no_topk→top_k,no_topp→top_p,no_topk_no_topp→both.sampling_batch_spec_dec_one_modelinNO_TOPKyields the exact same tokens asFULLwhen top_k is disabled, for thesame seed/offset, across
top_p ∈ {1.0, 0.9, 0.5}(the top_k mask is a no-opat
k=vocab). Skipped on CPU-only CI.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Summary by CodeRabbit
New Features
AdvancedSamplingModefor one-model MTP speculative decoding, exposed asDecodingBaseConfig.advanced_sampling_mode(defaultFULL) and plumbed viaSpecMetadata.FULL— apply both top-k and top-p filtersNO_TOPK— skip top-k filteringNO_TOPP— skip top-p renormalizationNO_TOPK_NO_TOPP— skip both filtersNone:resolve_advanced_sampling_filters(...)converts(advanced_sampling_mode, top_k, top_p)into effective optional tensors where disabled filters becomeNoneNonesample_from_logits_op(...):top_kis providedtop_pis providedtop_k/top_pfromSpecMetadata.advanced_sampling_modesample_from_logits_op(...)for non-greedy target token samplingTests
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.pycovering:skips_top_k/skips_top_pflagsDecodingBaseConfig/MTPDecodingConfigresolve_advanced_sampling_filters(...)identity vs disabled-to-NonebehaviorNO_TOPKmatchesFULLwhentop_kis disabled)FULLmodes across speculative/spec paths (including removal of the prior “one-model-only” gate)tensorrt_llm/usage/llm_args_golden_manifest.jsonto addspeculative_config.advanced_sampling_modeas a categorical argument.Docs / Config
docs/source/features/sampling.md, including semantics, deploy-time behavior (not in CUDA-graph key), independence from rejection sampling, and a usage example.Dev Engineer Review
Correctness & API consistency
resolve_advanced_sampling_filters(...)and consistently threaded into:compute_probs_from_logits)top_k/top_pasNonealigns with the stated goal: redundant FlashInfer kernels can be skipped while keeping mathematical no-op behavior when filters are disabled.CUDA-graph / keying
advanced_sampling_modeis excluded from the CUDA-graph key; functional selection is driven by resolved optional filter tensors (Nonevs provided tensors), matching the intended graph-stability requirement.Config / manifest
AdvancedSamplingModeadded totensorrt_llm/llmapi/llm_args.pyand defaulted toFULLonDecodingBaseConfig.tensorrt_llm/usage/llm_args_golden_manifest.jsonupdated to includespeculative_config.advanced_sampling_modewith allowed values:full,no_topk,no_topp,no_topk_no_topp.QA Engineer Review
Test changes
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.pyAdvancedSamplingModevalues and skip flags)DecodingBaseConfig,MTPDecodingConfig, independent of rejection sampling)resolve_advanced_sampling_filtersmaps disabled inputs toNone)no_topk/no_topk_no_topp)FULLmodes allowed and stored)tests/integration/test_lists/): not indicated in the provided context.